#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : tngs_select_chromosome.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/11/10
@Description : 从基因组文件中挑选出包含指定关键字的染色体序列
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
import re
from Bio import SeqIO


def tngs_select_chromosome(args: Namespace) -> None:
    """
    从下载的基因组文件中，根据关键字筛选完整染色体序列，并保存到输出文件。

    :param args: 命令行参数
    - in_dir: 输入目录，包含 .fna 文件
    - out_file: 输出 fasta 文件路径
    - key_word: 匹配描述行的关键字（正则支持）
    """
    # 清理关键字两端空格
    key_word = args.key_word.strip()

    # 输入文件夹
    in_dir = Path(args.in_dir).resolve()
    if not in_dir.exists() or not in_dir.is_dir():
        raise FileNotFoundError(f"输入目录不存在或不是文件夹: {in_dir}")

    # 输出文件
    out_file = Path(args.out_file).resolve()
    out_file.parent.mkdir(parents=True, exist_ok=True)  # 确保父目录存在

    # 收集所有 .fna 文件
    in_files = list(in_dir.glob("*.fna"))
    if not in_files:
        raise FileNotFoundError(f"输入目录中未找到 .fna 文件: {in_dir}")

    out_records = []

    # 遍历每个文件
    for in_file in in_files:
        # 打开 fasta 文件
        with open(in_file, "r", encoding="utf-8") as handle:
            for record in SeqIO.parse(handle, "fasta"):
                desc = record.description
                # 使用正则匹配关键字
                if re.search(key_word, desc):
                    # 修改 description 为 accession，保留序列原始信息
                    out_records.append(record)
                    continue

    # 如果没有匹配到任何序列，提醒用户
    if not out_records:
        print(f"未找到匹配关键字 '{key_word}' 的序列。")
        return

    # 写入输出文件
    with open(out_file, "w", encoding="utf-8") as out_handle:
        SeqIO.write(out_records, out_handle, "fasta")

    print(f"已筛选 {len(out_records)} 条序列，输出到: {out_file}")


if __name__ == "__main__":
    parser = ArgumentParser(description="从基因组文件中挑选出包含关键字的染色体序列")
    parser.add_argument("-i", "--in_dir", required=True, type=Path,
                        help="输入目录，包含 .fna 文件")
    parser.add_argument("-o", "--out_file", required=True, type=Path,
                        help="输出 fasta 文件路径")
    parser.add_argument("-k", "--key_word", required=True, type=str,
                        help="匹配描述行的关键字（正则支持）")
    arguments = parser.parse_args()
    tngs_select_chromosome(arguments)
